# -*- coding: utf-8 -*-
import numpy as np
import os, cv2, time, traceback
# from config_loader import *
from tools.utils import *
from tools.aug_paper import *
# from tools.ultralytics_predict import *
from tools.aug_defect_rgb import *
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
IS_GEN_ONLINE = False
INPUT_ORIGINAL_DATA_ROOT =    "D:/0_Data/_cd_data/1_ori_text/all"
OUTPUT_GEN_DEFECT_ROOT =      "D:/0_Data/_cd_data/2_gen_data"
OUTPUT_GEN_DEFECT_SUB_A =     "ng_test/A/"
OUTPUT_GEN_DEFECT_SUB_B =     "ng_test/B/"
OUTPUT_GEN_DEFECT_SUB_LABEL = "ng_test/label/"
DEFECT_MASK_ROOT =            "D:/0_Data/_cd_data/defect_mask"
SAMPLE_NUM = 2
IS_GEN_BACK = False
IMGSZ = 1024
DEFECT_SPECIAL_NAME = "RGB_final_text_1misspt"
PRINTING_DEFECT_NAME_DICT = {
                           # "text_dirty_pt_rgb": 0,
                           "text_stroke_miss_rgb": 0,
                           # "text_stroke_blur_rgb": 0,
                           # "back_color_miss_plus_rgb": 0,
                            # "back_ink_skin_rgb":0
                                }


def multi_thread_simulate(arg_list_):
    # 获取当前一个线程的数据，即arg_list_的其中一个列表元素
    (image_insp_path,
     image_temp_path,
     image_text_path,
     defect_mask_root,
     defect_name_dict,
     imgsz,
     image_id,
     ng_sim_root_a, ng_sim_root_b, ng_sim_root_label) = arg_list_

    # 根据参数获取数据
    image_insp = cv2.imread(image_insp_path, 1)
    image_temp = cv2.imread(image_temp_path, 1)
    image_text = cv2.imread(image_text_path, 0)
    image_h, image_w = image_insp.shape[0:2]

    # 在当前一个线程进行函数调用(gamma越小，图片越白)
    if random.random() > 0.5:
        image_insp_gamma = aug_light(image_insp, gamma=(0.75, 0.95))          # PDF: gamma=(0.4, 0.7)
    else:
        image_insp_gamma = aug_light(image_insp, gamma=(1, 1.3))          # PDF: gamma=(1.5, 1.75)
    image_insp_shift = aug_shift_defect(image_insp_gamma, shift_x=(-3, 3))  # Online +-4, Offline +-2
    image_text_shift = aug_shift_defect(image_text, shift_x=(-3, 3))        # Online +-4, Offline +-2

    # 创建一个包含所有启用的缺陷类型的列表
    enabled_defects_list = [name for name in defect_name_dict if name in defect_name_dict]
    defect_blob_mask = np.zeros((image_h, image_w), np.uint8)
    if not enabled_defects_list == []:
        defect_id = ""
        chosen_defect_name = random.choice(enabled_defects_list)
        if "back_ink_skin_rgb" == chosen_defect_name:
            defect_id = defect_name_dict["back_ink_skin_rgb"]
            image_aug, defect_blob_mask, is_gen_flag = aug_ink_skin_rgb(image_insp_shift,
                                                            defect_num=(5, 8),
                                                            thickness=(4, 10))
        elif "text_dirty_pt_rgb" == chosen_defect_name:
            defect_id = defect_name_dict["text_dirty_pt_rgb"]
            image_aug, defect_blob_mask, is_gen_flag = aug_dirty_pt_rgb(image_insp_shift,
                                                            image_text_shift,
                                                            is_gen_background=IS_GEN_BACK,
                                                            defect_type_list=['circle', 'ellipse', 'polygon'],
                                                            # defect_type_list=['line'],
                                                            defect_num=(5, 15),
                                                            pt_size=(3, 8),
                                                            line_size=(40, 200))  # (20, 400)
        elif "text_stroke_miss_rgb" == chosen_defect_name:
            defect_id = defect_name_dict["text_stroke_miss_rgb"]
            image_aug, defect_blob_mask, is_gen_flag = aug_text_miss_rgb(image_insp_shift,
                                                             image_text_shift,
                                                             ['circle', 'ellipse', 'polygon'],
                                                             defect_num=(5, 10),
                                                             pt_size=(5, 7))
        elif "text_stroke_blur_rgb" == chosen_defect_name:
            defect_id = defect_name_dict["text_stroke_blur_rgb"]
            image_aug, defect_blob_mask, is_gen_flag = aug_text_blur_rgb(image_insp_shift,
                                                                 image_text_shift,
                                                                 ['circle', 'ellipse', 'polygon'],
                                                                 defect_num=(10, 15),
                                                                 pt_size=(8, 15))
        elif "back_color_miss_plus_rgb" == chosen_defect_name:
            defect_id = defect_name_dict["back_color_miss_plus_rgb"]
            image_aug, defect_blob_mask, is_gen_flag = aug_color_miss_plus_rgb(image_insp_shift,
                                                                   defect_mask_root,
                                                                   defect_num=(3, 5),
                                                                   scale_range=(0.5, 3))  # 0.2, 2
        if is_gen_flag is False:
            return None
    else:
        image_aug = image_insp_shift.copy()


    # 确保图片必须按照指定的训练尺寸进行存储
    if image_h != imgsz or image_w != imgsz:
        image_aug = cv2.resize(image_aug, (imgsz, imgsz))
        image_temp = cv2.resize(image_temp, (imgsz, imgsz))
        defect_blob_mask = cv2.resize(defect_blob_mask, (imgsz, imgsz), interpolation=cv2.INTER_NEAREST)

    # 如果希望生成NG样本(只要在生成缺陷列表中存在缺陷关键字)
    if not enabled_defects_list == []:
        defect_name = get_key_from_value(defect_name_dict, defect_id)
        if (defect_name == "text_stroke_miss_rgb" or defect_name == "text_dirty_pt_rgb" or
                defect_name == "back_color_miss_plus_rgb" or defect_name == "text_stroke_blur_rgb" or
                defect_name=="back_ink_skin_rgb"):
            kernel = np.ones((9, 9), np.uint8)  # 将缺陷外扩一部分，便于生成OD的框更加大一点点(35, 35)。
            defect_blob_mask = cv2.dilate(defect_blob_mask, kernel, iterations=1)

    # --> 4) 各类仿真图像本地存储
    base_name = str(image_id)
    cv2.imwrite(ng_sim_root_a + "/" + DEFECT_SPECIAL_NAME + '_' + base_name + ".jpg", image_temp)
    cv2.imwrite(ng_sim_root_b + "/" + DEFECT_SPECIAL_NAME + '_' + base_name + ".jpg", image_aug)
    cv2.imwrite(ng_sim_root_label + "/" + DEFECT_SPECIAL_NAME + '_' + base_name + ".jpg", defect_blob_mask)
    return None

def gen_printing_defect_sample(imgsz,
                               is_gen_online,
                               defect_mask_root,
                               image_path_list_):
    # 建立文件夹
    defect_name_dict = PRINTING_DEFECT_NAME_DICT
    save_root = OUTPUT_GEN_DEFECT_ROOT
    ng_sim_root_a = os.path.join(save_root, OUTPUT_GEN_DEFECT_SUB_A)
    ng_sim_root_b = os.path.join(save_root, OUTPUT_GEN_DEFECT_SUB_B)
    ng_sim_root_label = os.path.join(save_root, OUTPUT_GEN_DEFECT_SUB_LABEL)

    # 自动创建指定的存储路径文件夹
    create_directory_if_not_exists(save_root)
    create_directory_if_not_exists(ng_sim_root_a)
    create_directory_if_not_exists(ng_sim_root_b)
    create_directory_if_not_exists(ng_sim_root_label)

    # 1) 通过多线程完成所有图片的仿真工作
    arg_list = []
    skip = 2 if is_gen_online is True else 3
    for image_id in range(0, len(image_path_list_), skip):
        image_insp_path = image_path_list_[image_id]
        if is_gen_online is True:
            image_temp_path = copy.deepcopy(image_insp_path)
            image_text_path = image_path_list_[image_id + 1]
        else:
            image_temp_path = image_path_list_[image_id + 1]
            image_text_path = image_path_list_[image_id + 2]
        arg_list.append((image_insp_path,
                         image_temp_path,
                         image_text_path,
                         defect_mask_root,
                         defect_name_dict,
                         imgsz,
                         image_id,
                         ng_sim_root_a, ng_sim_root_b, ng_sim_root_label))

    with ThreadPoolExecutor(max_workers=12) as executor:
        futures = [executor.submit(multi_thread_simulate, arg) for arg in arg_list]

        # 使用tqdm创建进度条
        for future in tqdm(as_completed(futures), total=len(futures), desc='Simulation progress:', ncols=100):
            try:
                pass
            except Exception as e:
                print("Error occurred: {e}")
    arg_list.clear()


def sim_printing_defect():
    """
    二维码缺陷仿真内容：
        基于定位ROI内对二维码实现缺陷仿真：模糊、缺印、断线、蹭脏
    """
    # 1) 获取原始OK图片数据路径
    image_path_list = get_file_path_list(INPUT_ORIGINAL_DATA_ROOT, FileType.IMAGE_TYPE)

    # 随机选择图片列表，不让数据量这么大
    skip = 2 if IS_GEN_ONLINE is True else 3
    rand_choic_num = len(image_path_list)//skip//SAMPLE_NUM
    new_image_path_list = sample_triplets(image_path_list, rand_choic_num, skip)

    # 2）缺陷仿真
    gen_printing_defect_sample(IMGSZ,
                               IS_GEN_ONLINE,
                               DEFECT_MASK_ROOT,
                               new_image_path_list)

if __name__ == '__main__':
    sim_printing_defect()


